import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from unitree_sdk2py.core.channel import ChannelFactoryInitialize
from unitree_sdk2py.go2.video.video_client import VideoClient
import cv2
import numpy as np
import sys

'''
主要改造点：
创建 ROS2 节点：将原有代码封装为 ImagePublisher 类继承自 Node
图像发布：
使用 sensor_msgs/msg/Image 消息类型
通过 cv_bridge 将 OpenCV 图像转换为 ROS2 图像消息
定时器回调：使用 create_timer 实现固定频率的图像获取和发布
消息头信息：添加了时间戳和坐标系信息
退出处理：完善了 ROS2 节点的关闭逻辑

下载AGIROS包及三方：
sudo apt install agiros-loong-cv-bridge agiros-loong-sensor-msgs
pip install numpy opencv-python

用法：
agiros run your_package go2_camera_publisher

查看图像话题：
agiros topic echo /camera/go2/image_raw
或使用 rqt_image_view 查看图像



'''
class ImagePublisher(Node):
    def __init__(self):
        super().__init__('go2_image_publisher')
        self.publisher_ = self.create_publisher(Image, 'camera/go2/image_raw', 10)
        self.bridge = CvBridge()
        self.timer = self.create_timer(0.05, self.timer_callback)  # 20Hz
        
        # Initialize video client
        if len(sys.argv) > 1:
            ChannelFactoryInitialize(0, sys.argv[1])
        else:
            ChannelFactoryInitialize(0)

        self.client = VideoClient()
        self.client.SetTimeout(3.0)
        self.client.Init()
        self.code = 0
        self.data = None

    def timer_callback(self):
        # Get image from Go2 robot
        self.code, self.data = self.client.GetImageSample()
        
        if self.code == 0:
            # Convert to numpy image
            image_data = np.frombuffer(bytes(self.data), dtype=np.uint8)
            cv_image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
            
            # Convert to ROS2 Image message
            ros_image = self.bridge.cv2_to_imgmsg(cv_image, encoding="bgr8")
            ros_image.header.stamp = self.get_clock().now().to_msg()
            ros_image.header.frame_id = "camera_link"
            
            # Publish the image
            self.publisher_.publish(ros_image)
            
            # Display image (optional)
            # cv2.imshow("front_camera", cv_image)
            if cv2.waitKey(1) == 27:  # ESC to stop
                self.destroy_node()
                rclpy.shutdown()


def main(args=None):
    rclpy.init(args=args)
    image_publisher = ImagePublisher()
    
    try:
        rclpy.spin(image_publisher)
    except KeyboardInterrupt:
        pass
    finally:
        image_publisher.destroy_node()
        rclpy.shutdown()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
